"""
========================================================================
FILE: geometric_field.py

LAST MODIFIED: 23 September 2009

DESCRIPTION: 
class for a ensemble_field_function representing a mesh geometry.
========================================================================
"""
"""
BEGIN LICENSE BLOCK 
Version: MPL 1.1/GPL 2.0/LGPL 2.1

The contents of this file are subject to the Mozilla Public License Version
1.1 (the "License"); you may not use this file except in compliance with
the License. You may obtain a copy of the License at
http://www.mozilla.org/MPL/

Software distributed under the License is distributed on an "AS IS" basis,
WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License
for the specific language governing rights and limitations under the
License.

The Original Code is fieldwork.

The Initial Developer of the Original Code is Ju Zhang.
Portions created by the Initial Developer are Copyright (C) 2009
the Initial Developer. All Rights Reserved.

Contributor(s):

Alternatively, the contents of this file may be used under the terms of
either the GNU General Public License Version 2 or later (the "GPL"), or
the GNU Lesser General Public License Version 2.1 or later (the "LGPL"),
in which case the provisions of the GPL or the LGPL are applicable instead
of those above. If you wish to allow use of your version of this file only
under the terms of either the GPL or the LGPL, and not to allow others to
use your version of this file under the terms of the MPL, indicate your
decision by deleting the provisions above and replace them with the notice
and other provisions required by the GPL or the LGPL. If you do not delete
the provisions above, a recipient may use your version of this file under
the terms of any one of the MPL, the GPL or the LGPL.

END LICENSE BLOCK
"""

import scipy
import template_fields
from enthought.mayavi import mlab

class geometric_field:
	""" Class for an ensemble_field representing the geometry of an 
	object. Either instantiated with an existing ensemble_field, or if 
	none, a new ensemble_field will be instantiated from the 
	template_fields module with the given field_dimension and basis_type.
	"""
	
	def __init__(self, name, dimensions, ensemble_field_function = None, field_dimensions = None, field_basis = None ):
		self.name = name
		self.dimensions = dimensions
		self.field_parameters = []
		self.points = []					# list of point objects
		self.points_to_ensemble_map = {}	# { points index: ensemble_point_number }
		self.ensemble_to_points_map = {}	# { ensemble_point_number: points index }
		self.named_points_map = {}			# { point_name: points index }
		self.points_counter = 0
		self.ensemble_point_counter = 0
		for i in range( self.dimensions ):
			self.field_parameters.append( [] )

		if ensemble_field_function:
			# create the appropriate number of geometric points for the ensemble
			# points in the mesh
			self.ensemble_field_function = ensemble_field_function
			self._create_ensemble_points()
		else:
			# instantiate a new ensemble_field_function
			self.ensemble_field_function = template_fields.empty_field( self.name+'field', field_dimensions, field_basis  )
		
	#==================================================================#
	def _create_ensemble_points( self ):
		""" Creates a geometric point for each ensemble point in the
		ensemble field function upon instantiation with an existing 
		ensemble_field_function
		"""
		
		# get total number of points need to make
		n_points = self.ensemble_field_function.get_number_of_ensemble_points()
		
		for i in range( n_points ):
			# make geometric_point object
			point = geometric_point( self.dimensions, ensemble_point_number = i )
			self.points.append( point )
			self.ensemble_to_points_map[i] = self.points_counter
			self.points_to_ensemble_map[self.points_counter] = i
			self.points_counter += 1
			
		return 1
	
	#==================================================================#	
	def get_number_of_points( self ):
		""" returns total number of geometric_points in the
		geometric_field
		"""
		
		return len(self.points)
		
	#==================================================================#		
	def set_field_parameters( self, field_parameters ):
		""" Sets the field parameters for the ensemble_field. A 
		geometric point is create at each location.
		
		field_parameters should be a list with self.dimension number of 
		sub lists containing the parameters from each dimension
		"""
		
		# checks
		if len( field_parameters ) != self.dimensions:
			print 'ERROR: geometric_field.set_parameters: parameters inconsistent with number of dimensions'
			return None
			
		p_lengths = [ len(p) for p in field_parameters ]
		n_ensemble_points = self.ensemble_field_function.get_number_of_ensemble_points()
		if (sum(p_lengths) / self.dimensions) != p_lengths[0]:
			print 'ERROR: geometric_field.set_field_parameters: parameter vector lengths not equal'
			return None
			
		if p_lengths[0] != n_ensemble_points:
			print 'ERROR: geometric_field.set_field_parameters: number of given parameters (', p_lengths[0], ') does not match number of field ensemble points (', n_ensemble_points, ')'
			return None

		# if nothings wrong so far
		#~ self.field_parameters =  scipy.array( field_parameters )
		self.field_parameters[:] =  field_parameters[:]
		
		# set geometric points associated with ensemble point with these parameters
		for i in range( n_ensemble_points ):
			fp = [ v[i] for v in self.field_parameters ]
			self.points[ self.ensemble_to_points_map[i] ].set_field_parameters( fp )
		
		return 1
	
	#==================================================================#
	def add_geometric_point( self, field_parameters, name = None ):
		""" Add a point to the geometric field. 
		
		field parameters length must = field dimensions
		returns new point number
		"""
		
		if len( field_parameters ) != self.dimensions:
			print 'ERROR: geometric_field.add_point: parameters length/ dimension mismatch'
			return
		
		self.points.append( geometric_point( self.dimensions, field_parameters ) )
		if name:
			self.named_points_map[name] = self.points_counter
			
		self.points_counter += 1
		
		return ( self.points_counter - 1 )
		
	#==================================================================#			
	def add_element( self, element, points, debug = 0 ):
		""" Add an element object to the field, connecting its element
		points to the geometric point numbers given.
		 
		element is an element or subfield object.
		points is a list of geometric point numbers in the order of the 
		element points of the element being added.
		"""
		
		# check for right number of points
		element_points = element.get_number_of_ensemble_points()
			
		if element_points != len( points ):
			print 'ERROR: geometric_field.add_element: number of points', len( points ), 'does not match number of element points', element_points
			return
		
		# element point counter
		ep = 0
		# add element to mesh
		e_number = self.ensemble_field_function.add_element( element )
		
		# connect element and update mapper
		for p in points:	
			# if point is an existing ensemble point
			if p in self.points_to_ensemble_map.keys():
				ensemble_point = self.points_to_ensemble_map[p]
				# get element point tuple for current connection point
				connected_eps = self.ensemble_field_function.mapper.get_ensemble_point_element_points( ensemble_point )
				if debug:
					print 'connecting', (e_number, ep), 'to:', connected_eps
				
				# connect
				connected_eps.append( ( e_number, ep ) )
				self.ensemble_field_function.connect_element_points( connected_eps )
			else:
				# otherwise point maps to a new ensemble point	
				#~ new_ensemble_point = max( self.ensemble_to_points_map.keys() ) + 1
				self.points_to_ensemble_map[p] = self.ensemble_point_counter
				self.ensemble_to_points_map[ self.ensemble_point_counter ] = p
				self.ensemble_point_counter += 1
				
				# add new field parameters to field_parameters
				fp = self.points[p].get_field_parameters()
				for i in range( self.dimensions ):
					self.field_parameters[i].append( fp[i] )
			
			ep += 1
			
		# update mapper
		self.ensemble_field_function.map_parameters()
		
		return 1
	
	#==================================================================#
	def add_element_with_parameters( self, element, parameters ):
		""" Add an element to the ensemble_field with predefined parameters
		new geometric points will be created for new unique parameter
		coordinates. Non-uniques will be mapped to existing ensemble
		points.
		
		parameters should be [ [d0], [d1], [d2] ]
		[d0] = [ [p0], [p1], [p2] ] where [p0] is a list of the
		parameters are point 0
		"""

		# check parameters for right dimensionality and number of points
		element_points = element.get_number_of_ensemble_points()
		if scipy.array(parameters).shape[0:2] != ( self.dimensions, element_points ):
			print 'ERROR: geometric_field.add_element_with_parameters: parameters length/ dimension mismatch. need:', ( self.dimensions, element_points )
			return None
		
		# get geometric point numbers for connecting the element to
		existing_positions = self.get_all_point_positions()
		connect_points = []
		for i in range( element_points ):
			# ith parameter coordinates
			coord = [ p[i][0] for p in parameters ]
			try: 
				connect_points.append( existing_positions.index( coord ) )
			except ValueError:
				# create new geometric point at current parameter 
				# coordinates, and record new geometric point number
				# for connection
				connect_points.append( self.add_geometric_point( [ p[i] for p in parameters ] ) )
				
		# add element
		self.add_element( element, connect_points )
		
		return 1

	#==================================================================#
	#~ def interactive_add_element( self, element, debug = 0):
		#~ # add an existing element object to the geometric field using an
		#~ # interactive mayavi interface
		#~ 
		#~ # get number of points needed
		#~ element_points = element.get_number_of_ensemble_points()
		#~ 
		#~ # generate the scene in which the point picker will operate
		#~ field_eval_density = 10
		#~ f = mlab.figure(1)
		#~ self._plot_points( label = True )
		#~ self._plot_field( field_eval_density )
		#~ mlab.show
		#~ 
		#~ # get the points to which the element will connect to using the
		#~ # mayavi interactive points picker
		#~ 

		
		 			
	#==================================================================#
	def get_all_point_positions( self ):
		""" Returns a list containing a self.dimensions long list of
		coordinates for each geometric point
		"""
		
		positions = []
		
		for i in range( self.get_number_of_points() ):
			positions.append( self.get_point_position( i ) )
				
		return positions
	
	#==================================================================#		
	def get_point_position( self, point ):
		""" Returns the coordinates of a point
		"""
		
		position = []
		try:
			parameters = self.points[point].get_field_parameters()
		except IndexError:
			print 'ERROR: geometric_field.get_point_position: invalid point number'
			return None
		else:
			for i in range(self.dimensions):
				position.append( parameters[i][0] )
			
			return position	
	
	#==================================================================#
	def evaluate_geometric_field( self, density ):
		""" evaluates the field for all parameter components.
		Returns a list of self.dimension lists
		"""
		evaluation = []
		
		# evaluate each coordinate in field
		for d in self.field_parameters:
			evaluation.append( scipy.array(self.ensemble_field_function.evaluate_field_in_mesh( density, d )).flatten() )
		
		return evaluation
					
	#==================================================================#
	def display_geometric_field( self, field_eval_density, point_glyph = 'sphere', point_label = None, field_glyph = 'point', field_scale = 1.0 ):
		""" Plots all geometric points and point evaluated over the field
		in a mayavi scene.
		"""
		
		f = mlab.figure()
		mlab.clf()
		self._plot_points( point_glyph, label = point_label )
		self._plot_field( field_eval_density, field_glyph, field_scale )
		
		mlab.show()
		
		return 1
			
	#==================================================================#
	def _plot_points( self, glyph = 'sphere', label = None ):
		""" uses mayavi points3d to show the positions of all points 
		(with labels if label is true)
		
		label can be 'all', or 'landmarks'
		"""
		
		# get point positions
		p = scipy.array( self.get_all_point_positions() )
		
		if self.dimensions == 3:
			#~ f = mlab.figure()
			points_plot = mlab.points3d( p[:,0], p[:,1], p[:,2], mode = 'sphere', scale_factor = 2.0, color = (1,0,0) )
			
			# label all ensemble points with their index number
			if label == 'all':
				labels = range( len( self.points ) )
				for i in range( len(labels ) ):
					l = mlab.text( p[0,i], p[1,i], str(labels[i]), z = p[2,i], line_width = 0.01, width = 0.005*len(str(labels[i]))**1.1 )
	
			elif label == 'landmarks':
				m = self.named_points_map
				labels = m.keys()
				for label in labels:
					l = mlab.text( p[m[label]][0], p[m[label]][1], label, z = p[m[label]][2], line_width = 0.01, width = 0.005*len( label )**1.1 )
	
		return points_plot
	
	#==================================================================#
	def _plot_field( self, density, glyph = 'point', scale_factor = 1.0 ):
		""" uses mayavi points3d to show the evaluated field
		"""
		
		evaluation = []
		
		# evaluate each coordinate in field
		for d in self.field_parameters:
			evaluation.append( scipy.array(self.ensemble_field_function.evaluate_field_in_mesh( density, d )).flatten() )
			
		if self.dimensions == 3:
			field_plot = mlab.points3d( evaluation[0],evaluation[1],evaluation[2], mode = glyph, scale_factor = 0.8, color = (0.0,0.5,1.0))
			
		return field_plot
		

#======================================================================#
class geometric_point:
	""" Class of point objects used by geometric_field
	"""
	
	def __init__( self, dimensions, field_parameters = None, ensemble_point_number = None ):
		
		self.field_parameters = None
		self.ensemble_point_number = None

		self.dimensions = dimensions
		if field_parameters:
			if len(field_parameters) != dimensions:
				raise ValueError( 'ERROR: geometric_point.__init__: dimension and number of coordinates mismatch' )
				return
			else:
				self.field_parameters = tuple( field_parameters )
		if ensemble_point_number:
			self.ensemble_point_number = ensemble_point_number
			
		return
	
	#==================================================================#	
	def get_dimensions( self ):
		return self.dimensions
	
	#==================================================================#
	def set_field_parameters( self, field_parameters ):
		
		if len(field_parameters) != self.dimensions:
			raise ValueError( 'ERROR: geometric_point.set_field_parameters: dimension and number of coordinates mismatch' )
			return
			
		self.field_parameters = list( field_parameters )
		return 1
	
	#==================================================================#		
	def get_field_parameters( self ):
		return self.field_parameters
	
	#==================================================================#	
	def set_ensemble_point_number( self, n ):
		self.ensemble_point_number = n
	
	#==================================================================#
	def get_ensemble_point_number( self ):
		return self.ensemble_point_number

		
#=Fitting related======================================================#
#~ class geometric_field_fitter:
	#~ def __init__( self, geometric_field ):
		#~ 
		#~ self.geo_field = geometric_field
		#~ self.dimensions = self.geo_field.dimensions
	#~ 
	#~ def _get_element_map( self ):
		#~ # creates a mapping of all real elements to ensemble points
	#~ 
	#~ def _get_element_evaluators( self ):
		#~ # get element point evaluators for each element
			#~ 
	#~ def set_data( self, data ):
		#~ # data should be a list with self.dimension number of sub lists
		#~ # or arrays containing the coordinates of datapoints
		#~ # all sublists must have same length
		#~ 
		#~ if len( data ) != self.dimensions:
			#~ print 'ERROR: geometric_field.set_data: data inconsistent with number of dimensions'
			#~ return None
		#~ 
		#~ lengths = [ len(d) for d in data ]
		#~ if (sum(lengths) / self.dimensions) != lengths[0]:
			#~ print 'ERROR: geometric_field.set_data: data vector lengths not equal'
			#~ return None
			#~ 
		#~ self.fitting_data = data
		#~ 
		#~ return 1
		#~ 
	#~ def fit( self, iterations, search_elements ):
		#~ # iterations: number of fitting iterations
		#~ # search_elements: number of closest elements to perform closes distance search for each datapoint
		#~ 
		#~ # assign data points to elements
		#~ self.allocate_element()
		#~ 
		#~ # calculate initial error
		#~ self.calculate_error()
		#~ 
		
